package naivenoisy;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileReader;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.LinkedList;

import common.CorpusProcessorNoisy;
import common.Probability;

public class PosTaggerNaiveNoisy {

	private static HashMap<String, Integer> tCount;
	private static int tTotalCount;
	private static HashMap<String, HashMap<String, Integer>> ttCount;
	private static HashMap<String, HashMap<String, Integer>> wtCount;
	private static HashMap<String, String> mostFrequentSubsequentPOS;
	private static LinkedList<Integer> currentOptimalPosSequence;
	private static Probability currentMaxProb;
	private static ArrayList<ArrayList<String>> currentSentencePosGraph;
	private static ArrayList<String> currentSentence;
	private static int[] correctTags;

	public static void main(String[] args) {
		if (args.length != 1){
			System.out.println("Invalid argument, please use this format:");
			System.out.println("arguments: [n]");
			System.out.println("\tn = maximum sentence length");
			System.exit(0);
		}
		int n = Integer.valueOf(args[0]);
		System.out.println("Applying the naive noisy channel model tagger to the development set.");
		System.out.println("The entered value of n="+n+". Please note that n>9 could take several minutes...");
		tCount = new HashMap<String, Integer>();
		ttCount = new HashMap<String, HashMap<String, Integer>>();
		wtCount = new HashMap<String, HashMap<String, Integer>>();
		mostFrequentSubsequentPOS = new HashMap<String, String>();
		correctTags = new int[2];

		tTotalCount = CorpusProcessorNoisy.run(tCount, ttCount, wtCount,
				mostFrequentSubsequentPOS);

		try {
			File file = new File("CoNLL2009-ST-English-development-pos.txt");
			currentSentence = new ArrayList<String>();
			ArrayList<String> currentSentencePOS = new ArrayList<String>();
			int currentWordCount = 0;
			BufferedReader br = new BufferedReader(new FileReader(file));
			String line = null;
			
			while ((line = br.readLine()) != null) {
				if (!line.isEmpty()) {
					String[] cols = line.split("\t");
					currentWordCount = Integer.valueOf(cols[0]);
					if (currentWordCount <= n) {
						currentSentence.add(cols[2]);
						currentSentencePOS.add(cols[4]);
					}
				} else {
					if (currentWordCount <= n) {
						ArrayList<String> taggedPOS = tag();
						for (int i = 0; i < currentSentencePOS.size(); i++) {
							if (currentSentencePOS.get(i).equals(
									taggedPOS.get(i))) {
								correctTags[0]++;
							} else {
								correctTags[1]++;
							}
						}
						currentSentence = new ArrayList<String>();
						currentSentencePOS = new ArrayList<String>();
					} else {
						currentSentence = new ArrayList<String>();
						currentSentencePOS = new ArrayList<String>();
					}
				}

			}
			double acc = (double) correctTags[0]
					/ (correctTags[0] + correctTags[1]);
			System.out.println("Done tagging. Accuracy was "+ acc);
		} catch (IOException e) {
			e.printStackTrace();
		}
	}

	private static ArrayList<String> tag() {

		currentSentencePosGraph = getPosGraph(currentSentence);
		int[] possibilities = new int[currentSentence.size()];
		for (int i = 0; i < currentSentencePosGraph.size(); i++) {
			possibilities[i] = currentSentencePosGraph.get(i).size();
		}
		getOptimalSequence(possibilities);

		ArrayList<String> POStags = new ArrayList<String>();
		for (int i = 0; i < currentOptimalPosSequence.size(); i++) {
			int posIndex = currentOptimalPosSequence.get(i);
			POStags.add(currentSentencePosGraph.get(i).get(posIndex));
		}
		
		return POStags;
	}

	private static ArrayList<ArrayList<String>> getPosGraph(
			ArrayList<String> words) {
		ArrayList<ArrayList<String>> result = new ArrayList<ArrayList<String>>();
		for (int i = 0; i < words.size(); i++) {
			ArrayList<String> a = new ArrayList<String>();
			if (wtCount.containsKey(words.get(i))) {
				a = new ArrayList<String>(wtCount.get(words.get(i)).keySet());
			} else if (i == 0) {
				a.add(mostFrequentSubsequentPOS.get("<s>"));
			} else {
				for (int j = 0; j < result.get(i - 1).size(); j++) {
					a.add(mostFrequentSubsequentPOS.get(result.get(i - 1)
							.get(j)));
				}
			}
			result.add(a);
		}
		return result;
	}

	private static void getOptimalSequence(int[] possibilities) {
		currentMaxProb = new Probability();
		nextSequence(possibilities, possibilities.length,
				new LinkedList<Integer>());
	}

	private static void nextSequence(int[] possibilities, int n,
			LinkedList<Integer> sequence) {
		if (n != 0) {
			for (int i = 0; i < possibilities[n - 1]; i++) {
				sequence.push(i);
				nextSequence(possibilities, n - 1, sequence);
				sequence.pop();
			}
		} else {
			Probability prob = new Probability((double) 1);
			String prevPOS = "<s>";
			int prevPOSCount = tCount.get(prevPOS);
			for (int i = 0; i < sequence.size(); i++) {
				// P(ti | ti-1) term
				String currentPOS = currentSentencePosGraph.get(i).get(
						sequence.get(i));
				int currentPOSCount = tCount.get(currentPOS);
				Integer lastTwoPOSCount = ttCount.get(prevPOS).get(currentPOS);
				if (lastTwoPOSCount != null) {
					prob = prob.multiply(new Probability(
							(double) lastTwoPOSCount / prevPOSCount));
				} else {
					// Back-off to unigrams
					prob = prob.multiply(new Probability(
							(double) currentPOSCount / tTotalCount));
				}

				// P(wi | ti) term
				if (wtCount.containsKey(currentSentence.get(i))) {
					Integer currentWordPOScount = wtCount.get(
							currentSentence.get(i)).get(currentPOS);
					if (currentWordPOScount != null) {
						prob = prob
								.multiply(new Probability(
										(double) currentWordPOScount
												/ currentPOSCount));
					} else {
						// Back-off to unigrams
						prob = prob.multiply(new Probability(
								(double) currentPOSCount / tTotalCount));
					}
				} else {
					// Back-off to unigrams
					prob = prob.multiply(new Probability(
							(double) currentPOSCount / tTotalCount));
				}

				prevPOS = currentPOS;
				prevPOSCount = currentPOSCount;
			}
			if (prob.compareTo(currentMaxProb) == 1) {
				currentMaxProb = prob;
				currentOptimalPosSequence = new LinkedList<Integer>(sequence);
			}
		}
	}
}